// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <string_view>

#include <melon/container/f14_map.h>

#include <pollux/common/base/exceptions.h>

namespace kumo::pollux {
    class StringIdMap {
    public:
        static constexpr uint64_t kNoId = ~0UL;

        StringIdMap() = default;

        StringIdMap(const StringIdMap &other) = delete;

        StringIdMap(StringIdMap &&other) = delete;

        void operator=(const StringIdMap &other) = delete;

        void operator=(StringIdMap &&other) = delete;

        /// Returns the id of 'string' or kNoId if the string is not known.
        uint64_t id(std::string_view string);

        /// Returns the total length of strings involved in currently referenced
        /// mappings.
        int64_t pinnedSize() const {
            return pinnedSize_;
        }

        /// Returns the id for 'string' and increments its use count. Assigns a
        /// new id if none exists. Must be released with release() when no longer
        /// used.
        uint64_t makeId(std::string_view string);

        /// Returns the id for 'string' and increments its use count. Assigns a new id
        /// if none exists. Must be released with release() when no longer used.
        /// Recovers string id map by assigning 'id' to 'string' and increments its
        /// use count. The function returns the recovered 'id'. It throws if 'id' has
        /// already been assigned to other string or 'string' has already assigned a
        /// different id. Must be released with release() when no longer used.
        ///
        /// NOTE: this is used by SSD cache to recover the file name to file id
        /// mapping. This is to ensure the same file to be mapped to the same SSD file
        /// shard after recover.
        uint64_t recoverId(uint64_t id, std::string_view string);

        // Decrements the use count of id and may free the associated memory if no
        // uses remain.
        void release(uint64_t id);

        /// Increments the use count of 'id'.
        void addReference(uint64_t id);

        /// Returns a copy of the string associated with id or empty string if id has
        /// no string.
        std::string string(uint64_t id) {
            std::lock_guard<std::mutex> l(mutex_);
            auto it = idToEntry_.find(id);
            return it == idToEntry_.end() ? "" : it->second.string;
        }

        /// Resets StringIdMap.
        void testingReset() {
            std::lock_guard<std::mutex> l(mutex_);
            stringToId_.clear();
            idToEntry_.clear();
            lastId_ = 0;
            pinnedSize_ = 0;
        }

        uint64_t testingLastId() const {
            std::lock_guard<std::mutex> l(mutex_);
            return lastId_;
        }

    private:
        struct Entry {
            std::string string;
            uint64_t id;
            uint32_t numInUse{};
        };

        mutable std::mutex mutex_;
        melon::F14FastMap<std::string, uint64_t> stringToId_;
        melon::F14FastMap<uint64_t, Entry> idToEntry_;
        uint64_t lastId_{0};
        uint64_t pinnedSize_{0};
    };

    /// Keeps a string-id association live for the duration of this.
    class StringIdLease {
    public:
        StringIdLease() = default;

        /// Makes a lease for 'string' and makes sure it has an id.
        StringIdLease(StringIdMap &ids, std::string_view string)
            : ids_(&ids), id_(ids_->makeId(string)) {
        }

        StringIdLease(StringIdMap &ids, uint64_t id, std::string_view string)
            : ids_(&ids), id_(ids_->recoverId(id, string)) {
        }

        /// Makes a new lease for an id that already references a string.
        StringIdLease(StringIdMap &ids, uint64_t id) : ids_(&ids), id_(id) {
            ids_->addReference(id_);
        }

        StringIdLease(const StringIdLease &other) {
            ids_ = other.ids_;
            id_ = other.id_;
            if (ids_ && id_ != StringIdMap::kNoId) {
                ids_->addReference(id_);
            }
        }

        StringIdLease(StringIdLease &&other) noexcept {
            ids_ = other.ids_;
            id_ = other.id_;
            other.ids_ = nullptr;
            other.id_ = StringIdMap::kNoId;
        }

        void operator=(const StringIdLease &other) {
            clear();
            ids_ = other.ids_;
            if (ids_ && other.id_ != StringIdMap::kNoId) {
                ids_->addReference(other.id_);
            }
            id_ = other.id_;
        }

        void operator=(StringIdLease &&other) noexcept {
            clear();
            ids_ = other.ids_;
            id_ = other.id_;
            other.ids_ = nullptr;
            other.id_ = StringIdMap::kNoId;
        }

        ~StringIdLease() {
            clear();
        }

        void clear() {
            if (ids_) {
                ids_->release(id_);
                ids_ = nullptr;
                id_ = StringIdMap::kNoId;
            }
        }

        bool hasValue() const {
            return id_ != StringIdMap::kNoId;
        }

        uint64_t id() const {
            return id_;
        }

    private:
        StringIdMap *ids_{nullptr};
        uint64_t id_{StringIdMap::kNoId};
    };
} // namespace kumo::pollux
